package com.yxc.elasticsearchdemo.service;

import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.IOException;
import java.text.SimpleDateFormat;
import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author: YXC
 * @time: 2022/10/24 11:03
 **/
@Service
public class DownloadFromElasticsearchService {

    @Autowired
    private RestHighLevelClient client;

    private DateTimeFormatter formatter =  DateTimeFormatter.ofPattern("yyyy-MM-dd-HH");

    static Map<String, Integer> staticMap = new HashMap<>();

    static Integer count = 0;

    public String downloadByParams(long startTime, long endTime) throws IOException {
        System.out.println("开始查询");
        SearchRequest searchRequest = new SearchRequest("t_pkt_dns_alias");
        searchRequest.scroll(TimeValue.timeValueMinutes(1))
                .searchType(SearchType.DEFAULT)
                .source(
                        new SearchSourceBuilder()
                                .query(
                                        new RangeQueryBuilder("c_recvtime")
                                                .gte(startTime)
                                                .lte(endTime)
                                )
                                .size(10000)
                );
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        if (searchResponse.getHits().getHits().length > 0) {
            // 处理
            List<Map<String, Object>> dataList = new ArrayList<>();
            for (SearchHit hit : searchResponse.getHits().getHits()) {
                dataList.add(hit.getSourceAsMap());
            }

            staticFile(dataList);
        }
        String scrollId = searchResponse.getScrollId();
        while (StringUtils.isNotEmpty(scrollId)) {
            SearchScrollRequest searchScrollRequest =
                    new SearchScrollRequest(scrollId)
                            .scroll(TimeValue.timeValueMinutes(1));
            searchResponse = client.scroll(searchScrollRequest, RequestOptions.DEFAULT);
            scrollId = searchResponse.getScrollId();
            if (searchResponse.getHits().getHits().length > 0) {
                // 处理
                List<Map<String, Object>> dataList = new ArrayList<>();
                for (SearchHit hit : searchResponse.getHits().getHits()) {
                    dataList.add(hit.getSourceAsMap());
                }
                staticFile(dataList);
            }
        }
        System.out.println("写入文件中");
        File file = new File("/data/yxcTest");
//        File file = new File("D:\\data\\testData");
        if (!file.getParentFile().exists()) {
            file.getParentFile().mkdirs();
        }
        if (!file.exists()) {
            file.createNewFile();
        }
        List<String> dataList = new ArrayList<>();
        for (Map.Entry<String, Integer> entry : staticMap.entrySet()) {
            dataList.add(entry.getKey() + "\t" + entry.getValue());
        }
        FileUtils.writeLines(file, dataList);
        System.out.println("写入文件完成  /data/yxcTest");
        return "";
    }

    public void staticFile(List<Map<String, Object>> dataList) {
        for (Map<String, Object> dataMap : dataList) {
            String key = getKey(dataMap);
            if (staticMap.containsKey(key)) {
                staticMap.put(key, staticMap.get(key) + 1);
            } else {
                staticMap.put(key, 1);
            }
        }
        count = count + dataList.size();
        System.out.println("key：" + staticMap.keySet().size() + "  \t总数：" + count + "\t比例：" + ((double) staticMap.keySet().size()) / ((double) count));
    }

    private String getKey(Map<String, Object> dataMap) {
        StringBuffer sb = new StringBuffer();
        Long recvtimeLong = (Long) dataMap.get("c_recvtime");
        LocalDateTime receiveTime = Instant.ofEpochMilli(recvtimeLong).atZone(ZoneOffset.ofHours(8)).toLocalDateTime();

        sb.append(formatter.format(receiveTime)).append(":");
        sb.append(dataMap.get("c_enterprise_id")).append(":");
        sb.append(dataMap.get("c_src_ipv4")).append(":");
        sb.append(dataMap.get("c_sport")).append(":");
        sb.append(dataMap.get("c_dst_ipv4")).append(":");
        sb.append(dataMap.get("c_dport")).append(":");
        sb.append(dataMap.get("c_dns_ip")).append(":");
        sb.append(dataMap.get("c_dns_qry_name")).append(":");
        sb.append(dataMap.get("c_dns_resp_name")).append(":");
        sb.append(dataMap.get("c_d_boundary")).append(":");
        sb.append(dataMap.get("c_d_region")).append(":");
        sb.append(dataMap.get("c_hit_detail")).append(":");
        sb.append(dataMap.get("c_dns_intelligence_hit")).append(":");
        sb.append(dataMap.get("c_sip_intelligence_hit")).append(":");
        sb.append(dataMap.get("c_dip_intelligence_hit")).append(":");

        return sb.toString();
    }

}
